{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 5.3.2 NGSolve - PETSc interface\n",
    "We use the ngs2petsc interface to map vectors and matrices between NGSolve and PETSc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ipyparallel import Cluster\n",
    "c = await Cluster(engines=\"mpi\").start_and_connect(n=4, activate=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "from ngsolve import *\n",
    "from netgen.occ import unit_square\n",
    "comm = MPI.COMM_WORLD\n",
    "\n",
    "ngmesh = unit_square.GenerateMesh(maxh=0.1, comm=comm)\n",
    "    \n",
    "for l in range(2):\n",
    "    ngmesh.Refine()\n",
    "mesh = Mesh(ngmesh)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Python-module ngsolve.ngs2petsc provides functionality to transfer vectors and matrices between NGSolve and Python. \n",
    "\n",
    "Make sure that the ipyparallel server can import the module, e.g. by starting the cluster in the current directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "import ngsolve.ngs2petsc as n2p\n",
    "import petsc4py.PETSc as psc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "fes = H1(mesh, order=1, dirichlet=\"left|bottom\")\n",
    "u,v = fes.TnT()\n",
    "a = BilinearForm(grad(u)*grad(v)*dx+u*v*ds).Assemble()\n",
    "f = LinearForm(x*v*dx).Assemble()\n",
    "gfu = GridFunction(fes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The function CreatePETScMatrix takes an NGSolve matrix, and creates a PETSc matrix from it. A VectorMapping object can map vectors between NGSolve and PETSc. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "psc_mat = n2p.CreatePETScMatrix(a.mat, fes.FreeDofs())\n",
    "vecmap = n2p.VectorMapping (a.mat.row_pardofs, fes.FreeDofs())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create PETSc-vectors fitting to the matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "psc_f, psc_u = psc_mat.createVecs()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "setting up the parallel Krylov-space solver ...."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "ksp = psc.KSP()\n",
    "ksp.create()\n",
    "ksp.setOperators(psc_mat)\n",
    "ksp.setType(psc.KSP.Type.CG)\n",
    "ksp.setNormType(psc.KSP.NormType.NORM_NATURAL)\n",
    "ksp.getPC().setType(\"gamg\")\n",
    "ksp.setTolerances(rtol=1e-6, atol=0, divtol=1e16, max_it=400)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "moving vectors between NGSolve and PETSc, and solve:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "vecmap.N2P(f.vec, psc_f)\n",
    "ksp.solve(psc_f, psc_u)   \n",
    "vecmap.P2N(psc_u, gfu.vec);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gfu = c[:][\"gfu\"]\n",
    "from ngsolve.webgui import Draw"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "Draw (gfu[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### PETSc preconditioner for NGSolve\n",
    "Next we create a PETSc preconditioner, and wrap it into an NGSolve preconditioner:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "a = BilinearForm(grad(u)*grad(v)*dx+u*v*ds)\n",
    "# pre = Preconditioner(a, \"petsc\", pctype=\"gamg\", levels=10)\n",
    "pre = Preconditioner(a, \"gamg\")\n",
    "a.Assemble();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and use it in an NGSolve - CGSolver:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "from ngsolve.krylovspace import CGSolver\n",
    "inv = CGSolver(a.mat, pre, printrates=comm.rank==0)\n",
    "gfu.vec.data = inv * f.vec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gfu = c[:][\"gfu\"]\n",
    "Draw (gfu[0]);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
